Skip to main content

Find K Closest Elements

Question​

Given an array of sorted numbers and a target number, find the K closest elements to the target number in the array.

Example 1
Input: arr = [1,2,3,4,5], K = 4

Output: [2,3,4,5]

Solution​

all//Find K Closest Elements.py


def findKClosestElements(arr, target, k):
    n = len(arr)

    # Corner cases
    if target <= arr[0]:
        return arr[:k]
    if target >= arr[n - 1]:
        return arr[n - k:]

    # Doing binary search
    left = 0
    right = n - 1
    while left < right:
        mid = left + (right - left) // 2
        if arr[mid] == target:
            break
        if arr[mid] < target:
            left = mid + 1
        else:
            right = mid

    # Get the index of the element just smaller
    # than the target element
    if arr[mid] < target:
        mid += 1

    left = mid - 1
    right = mid
    count = 0

    # Expand on both sides till 'k' elements are
    # found or either of the pointer goes out of the
    # array bounds
    while left >= 0 and right < n and count < k:
        if target - arr[left] <= arr[right] - target:
            left -= 1
        else:
            right += 1
        count += 1

    # If k is more than the number of elements
    # available in the array
    if left < 0:
        return arr[:k]
    if right > n - 1:
        return arr[n - k:]

    return arr[left + 1:right]

arr = [12, 16, 22, 30, 35, 39, 42,
       45, 48, 50, 53, 55, 56]
target = 35
k = 4

print(findKClosestElements(arr, target, k))